package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
)

func showBytes(bytes []byte) string {
	var str []byte
	for _, b := range bytes {
		if b >= 33 && b <= 126 && b != 92 {
			str = append(str, b)
		} else {
			str = append(str, fmt.Sprintf("\\%02X", b)...)
		}
	}
	return string(str)
}

func writeAllTCP(con *net.TCPConn, message []byte) error {
	defer con.CloseWrite()
	for len(message) > 0 {
		l, err := con.Write(message)
		if err != nil {
			return err
		}
		message = message[l:]
	}
	return nil
}

func handleTCP(client *net.TCPConn, message []byte) {
	defer client.Close()
	addrStr := client.RemoteAddr().String()
	log.Printf("Client %s connected\n", addrStr)
	var data [1024]byte
	n, err := io.ReadFull(client, data[:])
	if err != nil && err != io.ErrUnexpectedEOF {
		log.Printf("Failed to read: %v\n", err)
		return
	}
	log.Printf("From %s: %s\n", addrStr, showBytes(data[:n]))
	if err := writeAllTCP(client, message); err != nil {
		log.Printf("Failed to write: %v\n", err)
		return
	}
}

func main() {
	if len(os.Args) != 5 || os.Args[1] != "tcp" && os.Args[1] != "udp" || os.Args[2] != "server" && os.Args[2] != "client" {
		fmt.Fprintf(os.Stderr, "Usage: %s tcp|udp server|client <server address> <message>\n", os.Args[0])
		os.Exit(1)
	}
	if os.Args[1] == "tcp" {
		addr, err := net.ResolveTCPAddr("tcp", os.Args[3])
		if err != nil {
			panic(err)
		}
		if os.Args[2] == "server" {
			con, err := net.ListenTCP("tcp", addr)
			if err != nil {
				panic(err)
			}
			defer con.Close()
			for {
				client, err := con.AcceptTCP()
				if err != nil {
					log.Printf("Failed to accept: %v\n", err)
					continue
				}
				go handleTCP(client, []byte(os.Args[4]))
			}
		} else {
			con, err := net.DialTCP("tcp", nil, addr)
			if err != nil {
				panic(err)
			}
			defer con.Close()
			if err := writeAllTCP(con, []byte(os.Args[4])); err != nil {
				panic(err)
			}
			var data [1024]byte
			n, err := io.ReadFull(con, data[:])
			if err != nil && err != io.ErrUnexpectedEOF {
				panic(err)
			}
			fmt.Println(showBytes(data[:n]))
		}
	} else {
		addr, err := net.ResolveUDPAddr("udp", os.Args[3])
		if err != nil {
			panic(err)
		}
		if os.Args[2] == "server" {
			con, err := net.ListenUDP("udp", addr)
			if err != nil {
				panic(err)
			}
			defer con.Close()
			for {
				var data [1024]byte
				n, addr, err := con.ReadFromUDP(data[:])
				if err != nil {
					log.Printf("Failed to read: %v\n", err)
					continue
				}
				log.Printf("From %v: %s\n", addr, showBytes(data[:n]))
				_, err = con.WriteTo([]byte(os.Args[4]), addr)
				if err != nil {
					log.Printf("Failed to write: %v\n", err)
					continue
				}
			}
		} else {
			con, err := net.DialUDP("udp", nil, addr)
			if err != nil {
				panic(err)
			}
			defer con.Close()
			_, err = con.Write([]byte(os.Args[4]))
			if err != nil {
				panic(err)
			}
			var data [1024]byte
			n, _, err := con.ReadFromUDP(data[:])
			if err != nil {
				panic(err)
			}
			fmt.Println(showBytes(data[:n]))
		}
	}
}
